/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * openUBMC is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#ifndef MCPY_BUILTINS_ASYNC_H
#define MCPY_BUILTINS_ASYNC_H

#include "any_value.h"
#include "lambda.h"
#include <exception>
#include <functional>
#include <mc/exception.h>
#include <mc/future.h>
#include <mc/memory.h>
#include <mc/result.h>
#include <mc/variant.h>
#include <memory>

#include <boost/asio/io_context.hpp>
#include <boost/preprocessor/control/iif.hpp>
#include <boost/preprocessor/seq/size.hpp>

namespace mcpy {

template <typename T>
struct is_result : std::false_type {};

template <typename T>
struct is_result<mc::result<T>>
    : std::true_type {};

template <typename T>
constexpr bool is_result_v = is_result<std::remove_cv_t<T>>::value;

template <typename F>
constexpr bool is_future_v = mc::is_future_v<F>;

template <typename T>
using future = mc::future<T, mc::runtime::any_executor>;

template <typename T>
using promise = mc::promise<T, mc::runtime::any_executor>;

// ensure_future: 将 mc::result<T> 桥接到 future
template <typename T>
inline auto ensure_future(mc::result<T> result) {
    return result.as_future();
}

// 重载：已经是 future 则直接返回
template <typename T, typename Executor>
inline auto ensure_future(mc::future<T, Executor> fut) {
    return fut;
}

template <typename T, typename... Args>
mc::result<T> call_async(const mc::variant& func, Args&&... args) {
    mc::variant result_var = mcpy::call(func, std::forward<Args>(args)...);
    return mc::result<T>(result_var.template as<T>());
}

// 异步执行上下文
// 注意：上下文的内容必须能够被拷贝且拷贝开销足够小，因为在触发异步化时会拷贝一次上下文
struct async_context {
    // 最大同步递归深度限制配置，超过该配额强制异步化
    static constexpr int max_sync_recursion_depth = 32;

    int recursion_depth = 0; // 当前递归深度
    // TODO:: 未来可扩展，例如添加调用追踪、性能统计等
};
namespace detail {
class coroutine_ref;
}

class coroutine {
public:
    coroutine() : value_(0) {
    }

    bool is_child() const {
        return value_ < 0;
    }
    bool is_parent() const {
        return !is_child();
    }
    bool is_complete() const {
        return value_ == -1;
    }

private:
    friend class detail::coroutine_ref;
    int value_;
};

namespace detail {
class coroutine_ref {
public:
    coroutine_ref(coroutine& c) : value_(c.value_) {
    }

    coroutine_ref(coroutine* c) : value_(c->value_) {
    }

    coroutine_ref(const coroutine_ref&) = default;

    inline ~coroutine_ref() {
        if (!modified_) {
            value_ = -1;
        }
    }

    inline operator int() const {
        return value_;
    }

    int& set_value(int v, bool in_yield_block) {
        modified_       = true;
        break_          = false;
        continue_       = false;
        yield_          = false;
        in_yield_block_ = in_yield_block;
        return value_   = v;
    }

    inline bool is_break() {
        return break_;
    }

    inline bool is_continue() {
        return continue_;
    }

    inline void set_break(bool f) {
        break_ = f;
    }

    inline void set_continue(bool f) {
        continue_ = f;
    }

    inline void yield() {
        yield_ = true;
    }

    inline bool is_yield() const {
        return yield_;
    }

    inline bool in_yield_block() const {
        return in_yield_block_;
    }

    inline void set_in_yield_block(bool f) {
        in_yield_block_ = f;
        if (!f) {
            // 退出一个 yield 块后设置协程状态为完成，除非遇到下一个 yield 块重新设置状态，
            // 否则运行完剩余代码后应该立即退出协程
            value_ = -1;
        }
    }

    inline void step() {
        if (value_ == 0) {
            // 如果业务代码中没有调用 co_yield 又回滚到开头，则直接退出协程
            value_ = -1;
        }
    }

private:
    void operator=(const coroutine_ref&);

    int& value_;
    bool modified_       = false;
    bool break_          = false;
    bool continue_       = false;
    bool yield_          = false;
    bool in_yield_block_ = false;
};
} // namespace detail

// 异步协程状态机（包含协程状态与 promise/future）
template <typename T, typename Derived>
struct async_impl : mc::enable_shared_from_this<Derived>, coroutine {
    // Promise
    mcpy::promise<T> promise_;

    // Future
    mcpy::future<T> future_;

    // 错误和取消状态
    std::exception_ptr error_;
    bool               canceled_ = false;

    async_impl()
        : promise_(mc::make_promise<T>(mc::runtime::any_executor{})),
          future_(promise_.get_future()) {
    }

    async_impl(const async_impl& other) = default;
    virtual ~async_impl() {
    }

    template <typename U>
    void set_result_impl(U&& value) {
        using ResumeValue = typename Derived::resume_value_type;

        if constexpr (std::is_same_v<std::decay_t<U>, ResumeValue>) {
            // resume_value_type 作为参数的情况只有 co_reenter(this) 块外部才会出现
            // 这个场景我们只接收 __resume_value 参数的异常信息，其他情况不接收使用默认值
            if (value.has_exception()) {
                promise_.set_exception(value.get_exception());
            } else {
                set_result();
            }
        } else if constexpr (std::is_same_v<std::decay_t<U>, std::exception_ptr>) {
            promise_.set_exception(std::forward<U>(value));
        } else if constexpr (std::is_void_v<T>) {
            promise_.set_value();
        } else {
            promise_.set_value(std::forward<U>(value));
        }
    }

    template <typename... U>
    mc::result<T> set_result(U&&... value) {
        if constexpr (sizeof...(U) > 0) {
            static_assert(sizeof...(U) == 1, "set_result only supports one argument");
            set_result_impl(std::forward<U>(value)...);
        } else if constexpr (std::is_void_v<T>) {
            promise_.set_value();
        } else {
            // 需要返回值但是没有提供返回值，返回一个默认值
            // 这本应该不会出现的，除非异常情况到达协程 co_reenter(this) {} 块外部
            // 例如：co_reenter(this){} 块内异常到达块外部，协程状态机被销毁
            // 但是这种情况应该被视为协程执行失败，而不是返回默认值
            // 这里返回一个默认值是为了兼容性，但不应该出现这种情况，或者应该抛出异常？
            promise_.set_value(T{});
        }
        return get_result();
    }

    void set_error(std::exception_ptr e) {
        error_ = e;
    }

    mc::result<T> get_result() {
        return mc::result<T>(mcpy::future(future_));
    }

    template <typename F>
    mc::result<T> await(F&& awaitable, async_context& ctx, bool force_async = false) {
        return await_impl(std::forward<F>(awaitable), ctx, force_async);
    }

private:
    template <typename Awaitable>
    auto await_impl(Awaitable&& awaitable, async_context& ctx, bool force_async)
        -> std::enable_if_t<is_result_v<Awaitable>, mc::result<T>> {
        return std::visit([this, &ctx, force_async](auto&& v) {
            using arg_type = std::decay_t<decltype(v)>;
            if constexpr (is_future_v<arg_type>) {
                return await_impl(std::forward<decltype(v)>(v), ctx, force_async);
            } else if constexpr (std::is_same_v<arg_type, mc::error_ptr>) {
                MC_UNUSED(this);
                return do_resume(
                    std::make_exception_ptr(mc::detail::make_method_call_exception(v)),
                    ctx, force_async);
            } else {
                MC_UNUSED(this);
                return do_resume(std::forward<decltype(v)>(v), ctx, force_async);
            }
        }, awaitable.get_variant());
    }

    template <typename Future>
    auto await_impl(Future&& future, async_context& ctx, bool force_async)
        -> std::enable_if_t<is_future_v<Future>, mc::result<T>> {
        // 优化：如果已就绪，同步处理
        if (future.is_ready()) {
            if (!future.is_rejected()) {
                return do_resume(future.get(), ctx, force_async);
            } else {
                return do_resume(std::make_exception_ptr(
                                     future.get_exception()),
                                 ctx, force_async);
            }
        }

        auto self = this->shared_from_this();
        future.then([self, ctx = ctx, force_async](auto&& result) mutable {
            // 异步化时，重新设置调用栈
            ctx.recursion_depth = 0;
            self->do_resume(std::forward<decltype(result)>(result), ctx, force_async);
        }).catch_error([self, ctx, force_async](const mc::exception& e) mutable {
            // 异步化时，重新设置调用栈
            try {
                e.dynamic_rethrow_exception();
            } catch (...) {
                ctx.recursion_depth = 0;
                auto eptr           = std::current_exception();
                self->do_resume(std::move(eptr), ctx, force_async);
            }
        });
        return self->get_result();
    }

private:
    template <typename Value>
    mc::result<T> do_resume(Value&& value, async_context& ctx, bool force_async) {
        using ResumeValue = typename Derived::resume_value_type;

        ResumeValue resume_val;
        if constexpr (std::is_same_v<std::decay_t<Value>, std::exception_ptr>) {
            resume_val.set_exception(std::move(value));
        } else {
            resume_val.emplace(std::forward<Value>(value));
        }

        // 检查递归深度是否超限
        if (force_async || ctx.recursion_depth >= async_context::max_sync_recursion_depth) {
            auto self = this->shared_from_this();
            mc::runtime::post([self, val = std::move(resume_val), ctx_copy = ctx]() mutable {
                ctx_copy.recursion_depth = 0;
                self->do_execute(val, ctx_copy);
            });
            return self->get_result();
        }

        // 正常同步路径：递增深度
        ctx.recursion_depth++;
        return do_execute(resume_val, ctx);
    }

    template <typename Value>
    mc::result<T> do_execute(Value&& resume_val, async_context& ctx) {
        try {
            return static_cast<Derived*>(this)->execute(resume_val, ctx);
        } catch (...) {
            // 捕获异常并传递到调用链上层
            error_ = std::current_exception();
            promise_.set_exception(std::current_exception());
            return get_result();
        }
    }
};

// 异步函数基类
template <typename T, typename Derived, typename Signature>
struct async_base : mcpy::lambda_base<Derived, Signature> {
    async_base()                  = default;
    async_base(const async_base&) = default;

    template <typename... Args>
    auto execute(Args&&... args) {
        using impl_type = typename Derived::impl_type;
        auto impl       = mc::make_shared<impl_type>(std::forward<Args>(args)...);

        mcpy::async_context                   ctx;
        typename impl_type::resume_value_type __resume;
        return impl->execute(__resume, ctx);
    }
};

MC_API mc::result<void> sleep(int duration_ms);

} // namespace mcpy

#define co_label_cat(a, b)     a##b
#define co_label_cat3(a, b, c) a##b##c

#define co_reenter(c)                                         \
    for (mcpy::detail::coroutine_ref _coro_value = c;         \
         !_coro_value.is_yield(); _coro_value.step())         \
        switch (_coro_value)                                  \
        case -1:                                              \
            if (_coro_value) {                                \
                /* 添加 goto 避免编译器警告 lable 没有使用 */ \
                goto __terminate_coro;                        \
__terminate_coro:                                             \
                _coro_value.set_value(-1, false);             \
                goto __yield_coro;                            \
__yield_coro:                                                 \
                _coro_value.yield();                          \
                _coro_value.set_in_yield_block(false);        \
                continue;                                     \
            } else                                            \
            case 0:

#define co_yield_impl(n)                                                                           \
    /* yield 块内部的 co_break 或 co_continue 应该到这里退出  */                                   \
    if (_coro_value.is_break()) {                                                                  \
        co_label_cat(break_of_loop_, n) : _coro_value.set_break(false);                            \
        _coro_value.set_in_yield_block(false);                                                     \
        break;                                                                                     \
    } else if (_coro_value.is_continue()) {                                                        \
        co_label_cat(continue_of_loop_, n) : _coro_value.set_continue(false);                      \
        _coro_value.set_in_yield_block(false);                                                     \
        continue;                                                                                  \
    }                                                                                              \
    for (_coro_value.set_value(n, true);;)                                                         \
        if (_coro_value == 0) {                                                                    \
        case (n):;                                                                                 \
            /* 到这里说明当前 yield 块执行结束，设置协程状态为退出直到下一个 yield 重新设置状态 */ \
            _coro_value.set_value(-1, false);                                                      \
            break;                                                                                 \
        } else                                                                                     \
            switch (_coro_value ? 0 : 1)                                                           \
                for (;;)                                                                           \
                case -1:                                                                           \
                    /* 检测当前 yield 块内部是否有 co_break 或 co_continue  */                     \
                    if (_coro_value.is_break())                                                    \
                        goto co_label_cat(break_of_loop_, n);                                      \
                    else if (_coro_value.is_continue())                                            \
                        goto co_label_cat(continue_of_loop_, n);                                   \
                    else if (_coro_value)                                                          \
                        goto __terminate_coro;                                                     \
                    else                                                                           \
                        for (;;)                                                                   \
                        case 1:                                                                    \
                            /* 检测当前 yield 块内部是否有 co_break 或 co_continue  */             \
                            if (_coro_value.is_break())                                            \
                                goto co_label_cat(break_of_loop_, n);                              \
                            else if (_coro_value.is_continue())                                    \
                                goto co_label_cat(continue_of_loop_, n);                           \
                            else if (_coro_value)                                                  \
                                goto __yield_coro;                                                 \
                            else                                                                   \
                            case 0:

#define co_yield co_yield_impl(__LINE__)

// co_yield_return: 用于单行的 return await(...) 形式
// 是 co_yield 的简化版本，不需要处理 break/continue
#define co_yield_return_impl(n)               \
    for (_coro_value.set_value(n, false);;)   \
        if (_coro_value == 0) {               \
        case (n):;                            \
            _coro_value.set_value(-1, false); \
            break;                            \
        } else                                \
            return

#define co_yield_return co_yield_return_impl(__LINE__)

#define co_break                        \
    if (_coro_value.in_yield_block()) { \
        _coro_value.set_break(true);    \
    }                                   \
    break;

#define co_continue                                                                      \
    if (_coro_value.in_yield_block()) {                                                  \
        /* 在 co_yield 块内，需要设置 continue 标志并 break 跳出 co_yield 块， */        \
        /* 状态机会跳转到正确的 continue 位置 */                                         \
        /* 这是一个保护防止用户在非 co_yield  块内误用 co_continue 导致异常退出状态机 */ \
        _coro_value.set_continue(true);                                                  \
        break;                                                                           \
    } else {                                                                             \
        continue;                                                                        \
    }

#define co_try                                \
    for (_coro_value.set_value(n, false);;)   \
        if (_coro_value == 0) {               \
        case (n):;                            \
            _coro_value.set_value(-1, false); \
            break;                            \
        } else                                \
        case 0:                               \
            try

#define co_label_cat3(a, b, c) a##b##c

// 宏常量：控制 co_catch 是否清除协程状态
#define CO_CLEAR_STATE true
#define CO_KEEP_STATE  false

// co_catch: 用于所有 try 块的 catch 块，根据参数决定是否清除协程状态
// 参数：
//   try_id: try 块的 ID（如 0, 1, 2...）
//   clear_state: 是否清除协程状态（CO_CLEAR_STATE 或 CO_KEEP_STATE）
// 使用方式：
//   - 外层（配对 co_yield try）：try { ... } co_catch(0, CO_CLEAR_STATE);
//   - 内层（普通 try）：try { ... } co_catch(1, CO_KEEP_STATE);
#define co_catch_impl(try_id, clear_state)                                     \
    catch (...) {                                                              \
        co_label_cat(__exception_ptr_try_, try_id) = std::current_exception(); \
        if constexpr (clear_state) {                                           \
            _coro_value.set_value(-1, false);                                  \
        }                                                                      \
        goto co_label_cat3(__try_, try_id, _catch_check);                      \
    }
#define co_catch_keep_state(try_id, ...)                                       \
    catch (...) {                                                              \
        co_label_cat(__exception_ptr_try_, try_id) = std::current_exception(); \
        goto co_label_cat3(__try_, try_id, _catch_check);                      \
    }
#define co_catch(...)                                                      \
    BOOST_PP_IIF(BOOST_PP_GREATER(BOOST_PP_VARIADIC_SIZE(__VA_ARGS__), 1), \
                 co_catch_impl,                                            \
                 co_catch_keep_state)(__VA_ARGS__)

// co_except: 定义异常处理标签并检查是否有异常
// 参数：
//   try_id: try 块的 ID（如 0, 1, 2...）
// 使用方式：co_except(0) { ... } 或 co_except(0) try { ... } co_catch(parent_id);
#define co_except(try_id) \
    co_label_cat3(__try_, try_id, _catch_check) : if (co_label_cat(__exception_ptr_try_, try_id))

// co_except_else: 定义 else 块检查
// 参数：
//   try_id: try 块的 ID（如 0, 1, 2...）
// 使用方式：co_except_else(1) { ... } 或 co_except_else(1) try { ... } co_catch(parent_id);
#define co_except_else(try_id) \
    if (co_label_cat(__no_exception_try_, try_id))

// co_finally: 定义 finally 标签
// 参数：
//   try_id: try 块的 ID（如 0, 1, 2...）
// 使用方式：
//   - 非嵌套：co_finally(0)
//   - 嵌套：co_finally(1) try { ... } co_catch(0);
#define co_finally(try_id) \
    co_label_cat3(__try_, try_id, _finally) :

// co_check_exception: 检查 __resume_value 是否有异常，如果有则获取异常并跳转到 catch_check
// 参数：
//   try_id: try 块的 ID（如 0, 1, 2...）
// 使用方式：co_check_exception(0);
#define co_check_exception(try_id)                                                   \
    if (__resume_value.has_exception()) {                                            \
        co_label_cat(__exception_ptr_try_, try_id) = __resume_value.get_exception(); \
        goto co_label_cat3(__try_, try_id, _catch_check);                            \
    }

// co_rethrow: 重新抛出异常（用于 except 处理中）
// 参数：
//   try_id: try 块的 ID（如 0, 1, 2...）
// 使用方式：co_rethrow(0) catch (const XXX&) { ... }
#define co_rethrow(try_id)                                                  \
    try {                                                                   \
        std::rethrow_exception(co_label_cat(__exception_ptr_try_, try_id)); \
    }

#endif // MCPY_BUILTINS_ASYNC_H
